
#include <rthw.h>
#include <rtthread.h>

#include "../bsp/PRINT/my_printf.h"

#define HEAP_MAGIC 0x1ea0

/* 内存控制块 */
struct heap_mem
{
    /* magic and used flag */
    rt_uint16_t magic;
    rt_uint16_t used;

    rt_size_t next, prev;
};

/* 堆基值 */
static rt_uint8_t *heap_ptr;

/** the last entry, always unused! */
static struct heap_mem *heap_end;

#define MIN_SIZE 12

#define MIN_SIZE_ALIGNED     RT_ALIGN(MIN_SIZE, RT_ALIGN_SIZE)
#define SIZEOF_STRUCT_MEM    RT_ALIGN(sizeof(struct heap_mem), RT_ALIGN_SIZE)

static struct heap_mem *lfree;   /* pointer to the lowest free block */

static struct rt_semaphore heap_sem;

/* 对齐后的可用内存大小 */
static rt_size_t mem_size_aligned;

void debug_sem(){
    static int num;

    printf("--- %x type: 0x%x\r\n", num++, rt_object_get_type(&heap_sem.parent.parent));
}


static void plug_holes(struct heap_mem *mem)
{
    struct heap_mem *nmem;
    struct heap_mem *pmem;

    /* plug hole forward */
    nmem = (struct heap_mem *)&heap_ptr[mem->next];
    if (mem != nmem &&
        nmem->used == 0 &&
        (rt_uint8_t *)nmem != (rt_uint8_t *)heap_end)
    {
        /* if mem->next is unused and not end of heap_ptr,
         * combine mem and mem->next
         */
        if (lfree == nmem)
        {
            lfree = mem;
        }
        mem->next = nmem->next;
        ((struct heap_mem *)&heap_ptr[nmem->next])->prev = (rt_uint8_t *)mem - heap_ptr;
    }

    /* plug hole backward */
    pmem = (struct heap_mem *)&heap_ptr[mem->prev];
    if (pmem != mem && pmem->used == 0)
    {
        /* if mem->prev is unused, combine mem and mem->prev */
        if (lfree == mem)
        {
            lfree = pmem;
        }
        pmem->next = mem->next;
        ((struct heap_mem *)&heap_ptr[mem->next])->prev = (rt_uint8_t *)pmem - heap_ptr;
    }
}


/* 初始化堆内存，在系统初始时被调用，起止地址通过链接脚本.lds确定 */
void rt_system_heap_init(void *begin_addr, void *end_addr)
{
    struct heap_mem *mem;
    rt_ubase_t begin_align = RT_ALIGN((rt_ubase_t)begin_addr, RT_ALIGN_SIZE);
    rt_ubase_t end_align = RT_ALIGN_DOWN((rt_ubase_t)end_addr, RT_ALIGN_SIZE);

    if((end_align > (2 * SIZEOF_STRUCT_MEM)) &&
        ((end_align - 2 * SIZEOF_STRUCT_MEM) >= begin_align)){  // 最少要存两个heap_mem
        mem_size_aligned = end_align - begin_align - 2 * SIZEOF_STRUCT_MEM;
    }else{
        printf("mem init, error begin address 0x%x, and end address 0x%x\n",
                   (rt_uint32_t)begin_addr, (rt_uint32_t)end_addr);
        
        return;
    }

    heap_ptr = (rt_uint8_t *)begin_align;

    mem = (struct heap_mem *)heap_ptr;
    mem->magic = HEAP_MAGIC;
    mem->used = 0;
    mem->prev = 0;  // pre、next是相对于heap_ptr的偏移量
    mem->next = mem_size_aligned + SIZEOF_STRUCT_MEM;   // mem_size_aligned是减去两个heap_mem后的大小

    heap_end = (struct heap_mem *)&heap_ptr[mem->next];
    heap_end->magic = HEAP_MAGIC;
    heap_end->used = 1;
    heap_end->prev = mem->next;
    heap_end->next = mem->next;

    rt_sem_init(&heap_sem, "heap", 1, RT_IPC_FLAG_FIFO);   // 初始化内存管理专用信号量

    lfree = (struct heap_mem *)heap_ptr; // lfree记录最头部的可用内存块
}

void *rt_malloc(rt_size_t size){
    struct heap_mem *mem, *mem2;
    rt_size_t ptr, ptr2;

    if(size == 0){
        return RT_NULL;
    }

    size = RT_ALIGN(size, RT_ALIGN_SIZE);

    if(size > mem_size_aligned){
        printf("no memory %d\r\n", size);

        return RT_NULL;
    }

    if(size < MIN_SIZE_ALIGNED){
        size = MIN_SIZE_ALIGNED;
    }

    /* 信号量保护资源 */
    rt_sem_take(&heap_sem, RT_WAITING_FOREVER);

    /* ptr是空闲内存的偏移量，ptr < mem_size_aligned - size说明可能存在能分配的块 */
    for(ptr = (rt_uint8_t *)lfree - heap_ptr; ptr < mem_size_aligned - size; ptr = ((struct heap_mem *)&heap_ptr[ptr])->next){
        mem = (struct heap_mem *)&heap_ptr[ptr];

        /* mem->next - (ptr + SIZEOF_STRUCT_MEM)是当前内存块大小 */
        if((!mem->used) && mem->next - (ptr + SIZEOF_STRUCT_MEM) >= size){
            
            /* 如果该内存块足够大，将剩余部分初始化成一个空闲块 */
            if(mem->next - (ptr + SIZEOF_STRUCT_MEM) >= size + SIZEOF_STRUCT_MEM + MIN_SIZE_ALIGNED){
                ptr2 = ptr + SIZEOF_STRUCT_MEM + size;

                mem2 = (struct heap_mem *)&heap_ptr[ptr2];
                mem2->magic = HEAP_MAGIC;
                mem2->used = 0;
                mem2->prev = ptr;
                mem2->next = mem->next;

                mem->next = ptr2;
                mem->used = 1;

                if(mem2->next != mem_size_aligned + SIZEOF_STRUCT_MEM){
                    ((struct heap_mem *)&heap_ptr[mem2->next])->prev = ptr2;
                }      
            }else{
                mem->used = 1;
            }

            /* 寻找新的、最头部的空闲块 */
            if(mem == lfree){
                while(lfree->used && lfree != heap_end)
                    lfree = (struct heap_mem *)&heap_ptr[lfree->next];
            }

            rt_sem_release(&heap_sem);

            return (rt_uint8_t *)mem + SIZEOF_STRUCT_MEM;
        }
    }

    rt_sem_release(&heap_sem);

    return RT_NULL;
}

void rt_free(void *rmem){
    struct heap_mem *mem;

    if(rmem == RT_NULL){
        printf("addr is invalid %d\r\n", 0);

        return;
    }

    if((rt_uint8_t *)rmem < heap_ptr || (rt_uint8_t *)rmem >= (rt_uint8_t *)heap_end){
        printf("addr is invalid %d\r\n", -1);

        return;
    }

    /* 获取信号量 */
    rt_sem_take(&heap_sem, RT_WAITING_FOREVER);

    mem = (struct heap_mem *)((rt_uint8_t *)rmem - SIZEOF_STRUCT_MEM);
    if (!mem->used || mem->magic != HEAP_MAGIC)
    {
        printf("to free a bad data block: %c\n", ' ');
        printf("mem: 0x%08x, used flag: %d, magic code: 0x%04x\n", (rt_uint8_t *)mem, mem->used, mem->magic);
    }
    RT_ASSERT(mem->used);
    RT_ASSERT(mem->magic == HEAP_MAGIC);    // 针对上面的if

    /* 标记为未使用即可 */
    mem->used = 0;
    mem->magic = HEAP_MAGIC;

    if(mem < lfree){
        lfree = mem;
    }

    /* 用于合并前后未使用的内存块 */
    plug_holes(mem);

    rt_sem_release(&heap_sem);

}